"""
工作流快速测试脚本
验证基本功能是否正常工作
"""

import pandas as pd
import numpy as np
import warnings
warnings.filterwarnings('ignore')

def test_workflow():
    """测试工作流基本功能"""
    print("🧪 开始工作流功能测试...")
    
    try:
        # 导入工作流
        from load_prediction_workflow import create_workflow
        print("✅ 工作流导入成功")
        
        # 创建模拟数据函数
        def simple_mock_data(start_time: str, end_time: str) -> pd.DataFrame:
            start_dt = pd.to_datetime(start_time)
            end_dt = pd.to_datetime(end_time)
            timestamps = pd.date_range(start_dt, end_dt, freq='H')[:-1]
            n_points = len(timestamps)
            
            if n_points == 0:
                return pd.DataFrame(columns=['timestamp', 'swt', 'swh', 'real_load'])
            
            # 简单的模拟数据
            hours = np.array([ts.hour for ts in timestamps])
            real_load = 100 + 50 * np.sin(2 * np.pi * hours / 24) + np.random.normal(0, 5, n_points)
            swt = 25 + 5 * np.sin(2 * np.pi * hours / 24) + np.random.normal(0, 1, n_points)
            swh = 0.01 + 0.005 * np.random.rand(n_points)
            
            return pd.DataFrame({
                'timestamp': timestamps,
                'swt': swt,
                'swh': swh,
                'real_load': real_load
            })
        
        # 创建工作流
        workflow = create_workflow(simple_mock_data)
        print("✅ 工作流创建成功")
        
        # 测试数据获取 - 增加到一周数据确保足够的样本
        start_time = "2024-01-01 00:00:00"
        end_time = "2024-01-10 00:00:00"  # 9天数据，确保有足够样本进行序列创建
        
        processed_data, raw_data = workflow.fetch_and_process_data(start_time, end_time)
        print(f"✅ 数据处理成功: {len(raw_data)} 条记录")
        
        # 检查数据格式
        required_keys = ['load', 'env', 'hour', 'dayofweek']
        for key in required_keys:
            if key in processed_data:
                print(f"   {key}: {processed_data[key].shape}")
        
        # 测试序列创建
        required_samples = workflow.model_config.seq_len + workflow.model_config.pred_len  # 168 + 24 = 192
        if len(raw_data) >= required_samples:  # 需要足够的数据创建序列
            X, y = workflow.prepare_training_data(start_time, end_time)
            print(f"✅ 序列创建成功: X包含{len(X)}个特征, y形状{y.shape}")
        else:
            print(f"⚠️ 数据量不足({len(raw_data)}<{required_samples})，跳过序列创建测试")
        
        print("🎉 所有基本功能测试通过!")
        return True
        
    except ImportError as e:
        print(f"❌ 导入失败: {e}")
        print("💡 请确保所有依赖文件都在正确位置")
        return False
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        print("💡 请检查代码逻辑或数据格式")
        return False

def test_data_function_interface():
    """测试数据函数接口兼容性"""
    print("\n🔌 测试数据函数接口...")
    
    def valid_data_function(start_time: str, end_time: str) -> pd.DataFrame:
        """标准格式的数据函数"""
        timestamps = pd.date_range(start_time, end_time, freq='H')[:-1]
        n = len(timestamps)
        return pd.DataFrame({
            'timestamp': timestamps,
            'swt': np.random.normal(25, 5, n),
            'swh': np.random.normal(0.01, 0.002, n),
            'real_load': np.random.normal(100, 20, n)
        })
    
    def invalid_data_function(start_time: str, end_time: str) -> pd.DataFrame:
        """错误格式的数据函数"""
        return pd.DataFrame({
            'time': ['2024-01-01'],
            'temperature': [25],
            'load': [100]
        })
    
    try:
        from load_prediction_workflow import create_workflow
        
        # 测试有效函数
        workflow1 = create_workflow(valid_data_function)
        _, data1 = workflow1.fetch_and_process_data("2024-01-01 00:00:00", "2024-01-01 06:00:00")
        print("✅ 标准数据函数接口测试通过")
        
        # 测试无效函数
        workflow2 = create_workflow(invalid_data_function)
        try:
            workflow2.fetch_and_process_data("2024-01-01 00:00:00", "2024-01-01 06:00:00")
            print("❌ 应该检测到格式错误但未检测到")
        except ValueError:
            print("✅ 格式错误检测正常")
        
    except Exception as e:
        print(f"❌ 接口测试失败: {e}")

def main():
    """主测试函数"""
    print("🚀 空调负荷预测工作流测试")
    print("=" * 50)
    
    # 基本功能测试
    success = test_workflow()
    
    if success:
        # 接口兼容性测试
        test_data_function_interface()
        
        print("\n" + "=" * 50)
        print("✅ 工作流测试完成!")
        print("\n📋 测试总结:")
        print("   ✓ 模块导入正常")
        print("   ✓ 工作流创建正常")
        print("   ✓ 数据处理功能正常")
        print("   ✓ 接口兼容性正常")
        print("\n🎯 现在可以使用以下方式开始:")
        print("   1. 运行 python workflow_demo.py 查看完整演示")
        print("   2. 参考 工作流使用指南.md 进行集成")
        print("   3. 使用 quick_train_and_predict 函数快速开始")
    else:
        print("\n❌ 测试失败，请检查环境配置")

if __name__ == "__main__":
    main()
